#!/usr/bin/env python3
# -*- coding: utf-8 -*-
# Copyright (c) Facebook, Inc. and its affiliates.
import os
from pathlib import Path
import json
import numpy as np
import tqdm
from PIL import Image
import cv2
from panopticapi.utils import IdGenerator, save_json
import pycocotools.mask as mask_util
import glob
from scipy import signal


g_output_dir = Path('/root/ws/data/stpls3d/processed_dir/')
data_dir = Path('/root/ws/data/stpls3d/')
# sub_dirs = [str(v) for v in range(1,26)]
sub_dirs = [str(v) for v in range(1,6)]


def convert(in_img_file,in_mask_file, out_img_file,out_mask_file):
    in_img = Image.open(in_img_file)
    in_mask = Image.open(in_mask_file)
    # w2 = int(in_img.size[0]/2)
    # h2 = int(in_img.size[1]/2)
    # in_img = in_img.resize((w2, h2), resample=Image.BILINEAR)
    # in_mask = in_mask.resize((w2, h2), resample=Image.NEAREST)
    in_img = np.array(in_img)
    in_mask = np.array(in_mask)
    assert in_img.dtype == np.uint8
    assert in_mask.dtype == np.uint8
    in_mask.flags.writeable = True

    in_mask = np.delete(in_mask,3,axis=2)
    in_img = np.delete(in_img,3,axis=2)
    
    Image.fromarray(in_img).save(out_img_file)
    Image.fromarray(in_mask).save(out_mask_file)
    

### 第一步，生成MASK格式图像和train/val目录结构
def gen_train_images(cdir):
    raw_dir = cdir
    processed_dir = g_output_dir
    train_dir = processed_dir / "train"
    val_dir = processed_dir / "val"  
    processed_dir.mkdir(parents=True, exist_ok=True)
    train_dir.mkdir(parents=True, exist_ok=True)
    val_dir.mkdir(parents=True, exist_ok=True)
    
    raw_img_dir = raw_dir / "images"
    raw_mask_dir = raw_dir / "masks"
    
    index = 0
    for file in tqdm.tqdm(list(raw_img_dir.iterdir())):
        
        farr = file.name.split("@")
        if farr[-5] != "-45":
            continue
        
        output_dir = train_dir
        index += 1
        in_img_file = raw_img_dir / file.name
        in_mask_file = raw_mask_dir / file.name
        out_img_dir = output_dir / "images" 
        out_mask_dir = output_dir / "masks" 
        out_img_dir.mkdir(parents=True, exist_ok=True)
        out_mask_dir.mkdir(parents=True, exist_ok=True)
        out_img_file = out_img_dir / file.name
        out_mask_file = out_mask_dir / file.name
        convert(in_img_file,in_mask_file,out_img_file,out_mask_file)

def gen_train_images_all():
    for sb in sub_dirs:
        disss = data_dir/ "raw_train" / sb
        gen_train_images(disss)
        
from scipy.interpolate import interp1d 
### 第一步，生成MASK格式图像和train/val目录结构
def gen_val_im(cdir,out_img_dir,out_mask_dir):    
    for file in tqdm.tqdm(list(cdir.iterdir())):
        fname = file.name
        if not fname.endswith(".jpg"):
            continue
        imgf = file
        fname = fname[:-4]
        ins_npzf = fname + "_InstanceLabels.npz"
        seg_npzf = fname + "_SemanticLabels.npz"
        ins_npz = np.load(cdir / ins_npzf)['data'].astype(np.int16)
        seg_npz = np.load(cdir / seg_npzf)['data']
        
        has_building = 1 in np.unique(seg_npz)
        if not has_building:
            continue
        
        # ha,wa = seg_npz.shape
        # steps = [i / 511.0 for i in range(512)]
        # was = [int((wa-1) * t) for t in steps]
        # has = [int((ha-1) * t) for t in steps]
        # coords = [[h,w] for w in was for h in has]
        # coords = np.array(coords)
        # seg_npz = seg_npz[coords[:,0],coords[:,1]]
        # seg_npz = np.where(((seg_npz == 1) | (seg_npz == 17)),1,0)
        # seg_npz = seg_npz.reshape(512,512,-1).astype(np.uint8)
        
        # ins_npz = ins_npz[coords[:,0],coords[:,1]]
        # ins_npz = ins_npz.reshape(512,512,-1).astype(np.int16)
        # ins_npz_gg = ins_npz % 256
        # ins_npz_bb = ins_npz >> 8
        # mask_array = np.concatenate((seg_npz,ins_npz_gg.astype(np.uint8),ins_npz_bb.astype(np.uint8)),axis=2)
        # mask_image = Image.fromarray(mask_array).resize((512, 512), resample=Image.NEAREST)
        non_inds = np.where((seg_npz != 1) )
        ins_npz[non_inds] = 0
        seg_npz[non_inds] = 0
        
        ins_npz_gg = ins_npz % 256
        ins_npz_bb = ins_npz >> 8
        
        mask_array = np.zeros((seg_npz.shape[0],seg_npz.shape[1],3),dtype=np.uint8)
        mask_array[:,:,0] = seg_npz.astype(np.uint8)
        mask_array[:,:,1] = ins_npz_gg.astype(np.uint8)
        mask_array[:,:,2] = ins_npz_bb.astype(np.uint8)
        mask_image = Image.fromarray(mask_array).resize((512, 512), resample=Image.NEAREST)
   
        in_img = Image.open(imgf)
        in_img = in_img.resize((512, 512), resample=Image.BILINEAR)
        
        out_img_file = out_img_dir / (fname + ".png")
        out_mask_file = out_mask_dir / (fname + ".png")
        
        in_img.save(out_img_file)
        mask_image.save(out_mask_file)
        pass

def gen_val_im_all():   
    for sb in sub_dirs:
        disss = data_dir/ "raw_val" / sb
        
        val_dir = g_output_dir / "val" / sb
        img_dir = val_dir / "images"
        mask_dir = val_dir / "masks"  
        val_dir.mkdir(parents=True, exist_ok=True)
        img_dir.mkdir(parents=False, exist_ok=True)
        mask_dir.mkdir(parents=False, exist_ok=True)
        gen_val_im(disss,img_dir,mask_dir)

import shutil

def select_val_im_files():
    
    g_val_dir = g_output_dir / "val"
    g_val_img_dir = g_val_dir / "images"
    g_val_mask_dir = g_val_dir / "masks"
    g_val_img_dir.mkdir(parents=False, exist_ok=True)
    g_val_mask_dir.mkdir(parents=False, exist_ok=True)
    
    for sb in sub_dirs:
        sb_val_dir = g_output_dir / "val" / sb
        sb_img_dir = sb_val_dir / "images"
        sb_mask_dir = sb_val_dir / "masks"  

        img_inds = 0
        
        for file in tqdm.tqdm(list(sb_img_dir.iterdir())):
            img_inds += 1
            if img_inds % 7 != 0:
                continue
            
            fname = file.name
            new_file_img = g_val_img_dir / fname
            new_file_mask = g_val_mask_dir / fname
            
            old_file_img = sb_img_dir / fname
            old_file_mask = sb_mask_dir / fname
            
            assert old_file_img.exists()
            assert old_file_mask.exists()

            shutil.copy(old_file_img,new_file_img)
            shutil.copy(old_file_mask,new_file_mask)
            

### 第二步  生成 ade20k_instance_imgCatIds.json 和 ade20k_instance_catid_mapping.txt
def gen_ins_catids():
    
    image_id = 0
    images = []
    for name, dirname in [("train", "train"), ("val", "val")]:
    #for name, dirname in [("train", "train")]:
    
        processed_dir = data_dir / "processed_dir"
        tv_dir = processed_dir / dirname / "images"
        for file in tqdm.tqdm(list(tv_dir.iterdir())):
            image = {}
            image["file_name"] = file.name
            image["id"] = image_id
            image_id += 1
            image["width"] = 512
            image["height"] = 512
            images.append(image)
    
    categories = []
    cat = {}
    cat["id"] = 1
    cat["name"] = "building"
    categories.append(cat)
    
    instance_imgCatIds = {}
    instance_imgCatIds["images"] = images
    instance_imgCatIds["categories"] = categories
    
    instance_imgCatIds_file = processed_dir / "instance_imgCatIds.json"
    save_json(instance_imgCatIds,instance_imgCatIds_file)
    
    with open(processed_dir / "instance_catid_mapping.txt",'w') as f2:
        f2.write("Instacne100\tSceneParse150\tFullADE20K\n")
        f2.write("1\t1\t311")
        f2.flush()
        
    
    pass

##第三步  生成COCO格式JSON
def gen_coco_json():
    processed_dir = data_dir / "processed_dir"
    # for name, dirname in [("train", "train"), ("val", "val")]:
    for name, dirname in [("val", "val")]:

        image_dir = processed_dir/ dirname / "images"
        instance_dir = processed_dir/ dirname / "masks"

        # img_id = 0
        ann_id = 1

        # json
        # out_file = os.path.join(str(processed_dir), f"/instance_{name}.json")
        out_file = processed_dir/ f"instance_{name}.json"

        # json config
        instance_config_file = processed_dir / "instance_imgCatIds.json"
        with open(instance_config_file) as f:
            category_dict = json.load(f)["categories"]

        # load catid mapping
        # it is important to share category id for both instance and panoptic annotations
        mapping_file = processed_dir / "instance_catid_mapping.txt"
        with open(mapping_file) as f:
            map_id = {}
            for i, line in enumerate(f.readlines()):
                if i == 0:
                    continue
                ins_id, sem_id, _ = line.strip().split()
                # shift id by 1 because we want it to start from 0!
                # ignore_label becomes 255
                map_id[int(ins_id)] = int(sem_id)

        for cat in category_dict:
            cat["id"] = map_id[cat["id"]]

        filenames = sorted(glob.glob(os.path.join(image_dir, "*.png")))

        ann_dict = {}
        images = []
        annotations = []

        for idx, filename in enumerate(tqdm.tqdm(filenames)):
            image = {}
            image_id = os.path.basename(filename)[:-4]

            image["id"] = image_id
            image["file_name"] = os.path.basename(filename)

            original_format = np.array(Image.open(filename))
            image["width"] = original_format.shape[1]
            image["height"] = original_format.shape[0]

            images.append(image)

            filename_instance = os.path.join(str(instance_dir), image_id + ".png")
            ins_seg = np.asarray(Image.open(filename_instance))
            assert ins_seg.dtype == np.uint8

            # 标签R波段表示类别
            # GB波段表示ID
            # B*256 + G
            # because 0 is reserved as VOID label
            ##
            instance_cat_ids = ins_seg[..., 0]
            gg = np.array(ins_seg[...,2],dtype=np.int32) << 8
            bb = ins_seg[..., 1]
            instance_ins_ids = gg + bb
#np.array(ins_seg[...,2],dtype=np.int32) << 8 
            # process things
            for thing_id in np.unique(instance_ins_ids):
                if thing_id == 0:
                    continue
                mask = instance_ins_ids == thing_id
                instance_cat_id = np.unique(instance_cat_ids[mask])
                assert len(instance_cat_id) == 1

                anno = {}
                anno['id'] = ann_id
                anno['ins_id'] = int(thing_id)
                
                ann_id += 1
                anno['image_id'] = image['id']
                anno["iscrowd"] = int(0)
                anno["category_id"] = int(map_id[instance_cat_id[0]])

                inds = np.nonzero(mask)
                ymin, ymax = inds[0].min(), inds[0].max()
                xmin, xmax = inds[1].min(), inds[1].max()
                anno["bbox"] = [int(xmin), int(ymin), int(xmax - xmin + 1), int(ymax - ymin + 1)]
                # if xmax <= xmin or ymax <= ymin:
                #     continue
                rle = mask_util.encode(np.array(mask[:, :, None], order="F", dtype="uint8"))[0]
                rle["counts"] = rle["counts"].decode("utf-8")
                anno["segmentation"] = rle
                anno["area"] = int(mask_util.area(rle))
                annotations.append(anno)

        # save this
        ann_dict['images'] = images
        ann_dict['categories'] = category_dict
        ann_dict['annotations'] = annotations

        save_json(ann_dict, out_file)


if __name__ == "__main__":
    # gen_mask_all()

    
    
    # gen_val_im_all()
    # select_val_im_files()

    # gen_ins_catids()
    gen_coco_json()